package com.spring.gpt.controller;

import com.spring.gpt.entity.RequestData;
import com.zhipu.oapi.ClientV4;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import io.reactivex.Flowable;
import reactor.core.publisher.Flux;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.concurrent.atomic.AtomicReference;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;

@RestController
@RequestMapping("/chat")
@CrossOrigin
@Slf4j
public class ChatController {

    private static final String requestIdTemplate = "mycompany-%d";
    private static final ObjectMapper mapper = new ObjectMapper();

    private final String APIKEY;
    private final ClientV4 client;

    public ChatController(@Value("${spring.ai.openai.api-key}") String apiKey) {
        this.APIKEY = apiKey;
        this.client = new ClientV4.Builder(apiKey).build();
    }


    /**
     * 流式获取回答后，返回拼接后的字符串（智谱https://open.bigmodel.cn）
     */
    @PostMapping(value = "/zp1")
    public String callWithMessage(@RequestBody RequestData requestData) {
        String msg = requestData.getMsg();
        //业务id
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());
        //参数工具
        List<ChatTool> chatToolList = new ArrayList<>();
        ChatTool chatTool = new ChatTool();
        chatTool.setType("code_interpreter");
        chatToolList.add(chatTool);
        //携带消息体，这个实例 是 官方提供
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), msg);
        messages.add(chatMessage);
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)
                .stream(Boolean.TRUE)
                .invokeMethod(Constants.invokeMethod)
                .messages(messages)
                .tools(chatToolList)
                .toolChoice("auto")
                .requestId(requestId)
                .build();
        //使用流式
        ModelApiResponse sseModelApiResp = client.invokeModelApi(chatCompletionRequest);

        if (sseModelApiResp.isSuccess()) {
            AtomicBoolean isFirst = new AtomicBoolean(true);
            List<Choice> choices = new ArrayList<>();
            AtomicReference<ChatMessageAccumulator> lastAccumulator = new AtomicReference<>();
            mapStreamToAccumulator(sseModelApiResp.getFlowable())
                    .doOnNext(accumulator -> {
                        {
                            if (isFirst.getAndSet(false)) {
                                log.info("Response: ");
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getTool_calls() != null) {
                                String jsonString = mapper.writeValueAsString(accumulator.getDelta().getTool_calls());
                                log.info("tool_calls: {}", jsonString);
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getContent() != null) {
                                log.info(accumulator.getDelta().getContent());
                            }
                            choices.add(accumulator.getChoice());
                            lastAccumulator.set(accumulator);
                        }
                    })
                    .doOnComplete(System.out::println)
                    .lastElement()
                    .blockingGet();
            ChatMessageAccumulator chatMessageAccumulator = lastAccumulator.get();
            ModelData data = new ModelData();
            data.setChoices(choices);
            if (chatMessageAccumulator != null) {
                data.setUsage(chatMessageAccumulator.getUsage());
                data.setId(chatMessageAccumulator.getId());
                data.setCreated(chatMessageAccumulator.getCreated());
            }
            data.setRequestId(chatCompletionRequest.getRequestId());
            sseModelApiResp.setFlowable(null);// 打印前置空
            sseModelApiResp.setData(data);
        }
        //组合内容
        StringBuilder stringBuilder = new StringBuilder();
        List<Choice> choices = sseModelApiResp.getData().getChoices();
        for (Choice choice : choices) {
            stringBuilder.append(choice.getDelta().getContent());
        }
        String content = stringBuilder.toString();
        // log.info("质谱ai的回答是：{}",content);
        return content;
    }

    /**
     * 通过ModelApiResponse.getFlowable()获取流式数据，最后通过blockingGet()获取最终结果
     * System.out.print(accumulator.getDelta().getContent());  // 这句代码是返回的具体结果
     * 因为直接SSE传text,受结束符\n影响，可以使用base64传输
     */
    @PostMapping(value = "/zp2", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> testSseInvoke(@RequestBody RequestData requestData) {
        String msg = requestData.getMsg();
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), msg);
        messages.add(chatMessage);
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)
                .stream(Boolean.TRUE)
                .messages(messages)
                .requestId(requestId)
                .build();
        Flux<String> flux = Flux.create(emitter -> {
            ModelApiResponse sseModelApiResp = client.invokeModelApi(chatCompletionRequest);
            if (sseModelApiResp.isSuccess()) {
                AtomicBoolean isFirst = new AtomicBoolean(true);
                mapStreamToAccumulator(sseModelApiResp.getFlowable())
                        .doOnNext(accumulator -> {
//                            if (isFirst.getAndSet(false)) {
//                                String base64Msg = Base64.getEncoder().encodeToString("Response: ".getBytes(StandardCharsets.UTF_8));
//                                emitter.next(base64Msg);
//                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getTool_calls() != null) {
                                try {
                                    String jsonString = mapper.writeValueAsString(accumulator.getDelta().getTool_calls());
                                    String tool = "tool_calls: " + jsonString;
                                    String base64Msg = Base64.getEncoder().encodeToString(tool.getBytes(StandardCharsets.UTF_8));
                                    emitter.next(base64Msg);
                                } catch (Exception e) {
                                    String err = "Error converting tool_calls to JSON: " + e.getMessage();
                                    String base64Msg = Base64.getEncoder().encodeToString(err.getBytes(StandardCharsets.UTF_8));
                                    emitter.next(base64Msg);
                                }
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getContent() != null) {
                                String content = accumulator.getDelta().getContent();
                                String base64Msg = Base64.getEncoder().encodeToString(content.getBytes(StandardCharsets.UTF_8));
                                // 具体结果
                                // System.out.print(content);
                                emitter.next(base64Msg);
                            }
                        })
                        .doOnComplete(() -> {
                            emitter.next("\n");
                            emitter.complete();
                        })
                        .subscribe();
            } else {
                emitter.next("Error: " + sseModelApiResp.getError() + "\n");
                emitter.complete();
            }
        });
        return flux;
    }

    //流式操作
    public static Flowable<ChatMessageAccumulator> mapStreamToAccumulator(Flowable<ModelData> flowable) {
        return flowable
                .map(chunk -> {
                    Delta delta = chunk.getChoices().get(0).getDelta();
                    Choice choice = chunk.getChoices().get(0);
                    Usage usage = chunk.getUsage();
                    String id = chunk.getId();
                    long created = chunk.getCreated();
            return new ChatMessageAccumulator(delta, null, choice, usage, created, id);
        });
    }

    /**
     * 测试传输
     */
    @GetMapping(value = "/zp3", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> testSseInvoke3(@RequestParam("msg") String msg) {
        String base64Msg = Base64.getEncoder().encodeToString(msg.getBytes(java.nio.charset.StandardCharsets.UTF_8));
        String base64AdditionalMsg = Base64.getEncoder().encodeToString("这是一条附加消息\nThis is an additional message".getBytes(java.nio.charset.StandardCharsets.UTF_8));
        return Flux.just(
                base64Msg + "\n\n",
                 base64AdditionalMsg + "\n\n"
        );
    }

}
